#ifndef SOLVER_H
#define SOLVER_H

#include <vector>

#include <petscsys.h>

#include <petsctime.h>

#include <petscksp.h>
#include <petscdmda.h>

#include "options.h"

#include "initial_data.h"

#include "maths.h"
#include "tools.h"

class Solver {
public:
  // Constructor, allocate space for temporary vectors, solution vectors and for laplacian matrix
  Solver(integer Nx_, integer Ny_, bool adaptive_ts_, real lambda_, real domain_[2][2], real T_, std::string model_, MPI_Comm local_comm_, int local_color_);
  Solver(Options *opt, MPI_Comm local_comm_, int local_color_);
  PetscErrorCode init(integer Nx_, integer Ny_, bool adaptive_ts_, real lambda_, real domain_[2][2], real T_, std::string model_, MPI_Comm local_comm_, int local_color_);
  
  // Free memory
  ~Solver();
  
  void set_snapshots(bool take, integer num, real *times);
  
  // Set varaibles
  void set_name(const char * name_) { strcpy(name, name_); }
  
  // Deconstruct every preallocated vector and matrix
  PetscErrorCode deallocate();
  
  // Returns the rank of processor, given the global grid coordinates
  integer get_rank(const integer & j, const integer & i);
  
  // Conversion between global grid coordinate and inline coordinates
  void global_to_inline(integer & I, const integer & j, const integer & i);
  void inline_to_global(integer & j, integer & i, const integer & I);
  
  // Return inline coordinate of neibourhoods of inline coordinate C
  void get_nhbrds(integer & L, integer & R, integer & B, integer & T, const integer & C);
  
  // Return the coordinates (x,y) relative to cell centered index (j, i)
  void get_coords(scalar & y, scalar & x, const integer & j, const integer & i);
  
  // Provide initial data
  void provide_initial_data();
  PetscErrorCode initialize_data();
  
  // Construct the laplacian matrix of the correct dimensions
  PetscErrorCode laplacian_matrix();
  
  // Performs all timesteppings up to time T, takes care of initial data
  PetscErrorCode solve();
  // Perform a time step on U, G local and global
  PetscErrorCode step();
  
  // Replace content of U with the convective derivative
  PetscErrorCode convection();
  // Replace the content of U and G with the projection of U and G
  PetscErrorCode projection();
  
  // Retrive the content of phi_global and writes its gradient on U_y, U_y
  PetscErrorCode gradient();
  // Retrive the content of U_y and places its divercence to Div_U_global
  PetscErrorCode divergence();
  
  // Collect data to process zero and calls deallocate()
  PetscErrorCode create_to_zero();
  
  // Exporting methods
  PetscErrorCode save_solution(QuantityType quantity = VELOCITY, DataType data = CURL, PlotType plot = NONE);
  PetscErrorCode save_energy(PlotType plot = NONE);
  PetscErrorCode save_log();
  
  // Mesh informations given with the constructor
  integer               Nx, Ny;
  bool                  adaptive_ts;
  real                  lambda;
  real                  T;
  real                  dx, dy;
  real                  left, right, bottom, top;
  
  std::string           model;
  
  // Counter for number of steps performed
  integer               step_counter;
  
  // Snapshot variables
  bool                  take_snapshot;
  integer               last_snapshot;
  integer               num_snapshot;
  real                  snapshot_times[MAX_SNAPSHOTS];
  bool                  do_snapshot;
  
  // mesh info computed at init, step
  real                  dt, t;
  real                  sup, l2;
  real                  elapsed, elapsed_cpu;
  std::vector<real>     times, energy;
  real                  area;
  real                  domain[2][2];
  
  // Local communicator
  MPI_Comm              local_comm;
  int                   local_color;
  
  // Number of processor in each direction
  integer               proc_x, proc_y;
  
  // Global data of grid partition for each processor
  integer               *vec_corn_x, *vec_corn_y;
  integer               *vec_span_x, *vec_span_y;
  integer               *cumu_areas;
  
  // Size and rank for MPI
  integer               comm_size, loc_rank;
  
  // Inline local indices
  integer               istart, iend;
  
  // Local grid space
  integer               corn_x, corn_y;
  integer               span_x, span_y;
  integer               ghost_corn_x, ghost_corn_y;
  integer               ghost_span_x, ghost_span_y;
  
  // Handlers for collected solution
  Vec                   U_collect, G_collect;
  
  // Name
  char                  name[PETSC_MAX_PATH_LEN];
  
  // Initial data
  InitialData           * initial_data;
  
private:
  // Error handler
  PetscErrorCode        ierr;
  
  // Mesh info (solver, DMDA)
  KSP                   ksp;
  DM                    da;
  integer               dof;
  integer               stencil;
  
  // Laplacian matrix
  Mat                   A;
  
  // Global vectors for the solution
  Vec                   U_local, G_local;
  // Temporary local vecvtors
  Vec                   P_local, U_x_local, U_y_local, U_no_x_local, U_no_y_local, U_edg_L_local, U_edg_R_local;
  // Local vectors for the solution
  Vec                   U_global, G_global;
  // Temporary global vecvtors
  Vec                   P_global, U_x_global, U_y_global, U_no_x_global, U_no_y_global, U_edg_L_global, U_edg_R_global;
  
  // Inline vectors
  Vec                   phi, Div_U;
};

#endif //SOLVER_H
